from crewai import Agent
from conn import llms
from pydantic import BaseModel, Field
import json


class ResponseModel(BaseModel):
    description: str = Field(..., description="实体描述")

def gen_description(url: str):
    agent = Agent(
        role="生成实体描述",
        goal="根据实体名生成实体的中文描述",
        llm=llms.get_llm(),
        backstory="非常专业的内实体描述生成助手",
        verbose=True,
    )

    r = agent.kickoff(url, response_format=ResponseModel)
    return json.loads(r.raw)['description']



if __name__ == "__main__":
    res = gen_description('AI2')
    print(res)